from openai import OpenAI
import os
import json
# 加载环境变量
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())  # 读取本地 .env 文件，里面定义了 OPENAI_API_KEY

# 实例化 OpenAI 客户端
client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url=os.getenv("OPENAI_BASE_URL")
)


# 移除字典中的 None 值
def remove_none_values(data):
    if isinstance(data, dict):
        return {key: remove_none_values(value) for key, value in data.items() if value is not None}
    elif isinstance(data, list):
        return [remove_none_values(item) for item in data if item is not None]
    else:
        return data


# 打印参数。如果参数是有结构的（如字典或列表），则以格式化的 JSON 形式打印；
def print_json(data):
    if (isinstance(data, (list, dict))):
        print(json.dumps(data, indent=4, ensure_ascii=False))
    else:
        print(data)


# 对象序列化
def serialize_json(data):
    if hasattr(data, 'model_dump_json'):
        data = json.loads(data.model_dump_json())

    if (isinstance(data, (list, dict))):
        data = remove_none_values(data)
    return data


# 封装 openai 接口
def get_completion(prompt, model="gpt-3.5-turbo"):
    params = [{"role": "user", "content": prompt}]
    print("=====GPT请求=====")
    messages = serialize_json(params)
    print_json(messages)

    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,  # 模型输出的随机性，0 表示随机性最小
    )
    return response.choices[0].message.content
